from datetime import datetime
import rsa
from rich.console import Console  # 控制台模块
from os.path import exists as os_path_exists
from os import makedirs as os_makedirs
console = Console(color_system="auto")  # 创建控制台


def create_keys(key_name):  # 生成公钥和私钥
    '''
    key_name 公钥名和私钥名
    '''
    (pubkey, privkey) = rsa.newkeys(1024)
    pub = pubkey.save_pkcs1()
    if not os_path_exists('./key'):
        os_makedirs('./key')  # 检查文件夹路径是否存在，如不存在则创建
    try:
        with open('./key/public_%s.pem' % key_name, 'wb+')as f:
            f.write(pub)

        pri = privkey.save_pkcs1()
        with open('./key/private_%s.pem' % key_name, 'wb+')as f:
            f.write(pri)
        console.log('[green]Key created successfully!')
        console.print(
            '[bold blue]Please keep your public and private keys \n请保存好您的公钥和私钥')

    except:
        console.print_exception()


def Public_key_encryption(text, key_path):  # 用公钥加密
    '''
    text 需要加密的文本
    '''
    try:
        Time = str(datetime.today()).split('.')[
            0].replace(' ', '_').replace(':', '-')
        if not os_path_exists('./text'):
            os_makedirs('./text')  # 检查文件夹路径是否存在，如不存在则创建

        with open('%s' % key_path, 'rb') as publickfile:
            p = publickfile.read()
        pubkey = rsa.PublicKey.load_pkcs1(p)
        original_text = text.encode('utf8')
        crypt_text = rsa.encrypt(original_text, pubkey)
        with open(r'./text/Encryption_text_%s_%s.txt' % (key_path.split('\\')[-1].split('.')[0].split('_')[1], Time), 'wb')as f:
            f.write(crypt_text)  # 将加密文本保存
        console.log('[green]已将加密密文保存至./text/Encryption_text_%s_%s.txt' %
                    (key_path.split('\\')[-1].split('.')[0].split('_')[1], Time))
        console.print('[#ffb74d]{Encryption 加密密文}：%s' % crypt_text)
        return crypt_text  # 加密后的密文
    except:
        console.print_exception()


def Private_key_decryption(crypt_text_path, key_path):  # 用私钥解密
    '''
    crypt_text_path 加密后的密文的路径
    key_path 私钥绝对路径
    '''
    try:  # 根据输入的文本解析为路径再对目标路径内文本进行解密
        Time = str(datetime.today()).split('.')[
            0].replace(' ', '_').replace(':', '-')
        if not os_path_exists('./text'):
            os_makedirs('./text')  # 检查文件夹路径是否存在，如不存在则创建
        with open('%s' % key_path, 'rb') as privatefile:
            p = privatefile.read()
        privkey = rsa.PrivateKey.load_pkcs1(p)
        # 注意，这里如果结果是bytes类型，就需要进行decode()转化为str
        with open('%s' % crypt_text_path, 'rb')as f:
            crypt_text = f.read()  # 读取加密文本
        last_text = rsa.decrypt(crypt_text, privkey).decode()
        with open(r'./text/Decryption_text_%s' % (crypt_text_path.split('\\')[-1]), 'w', encoding='gbk')as f:
            f.write(last_text)  # 将解密文本保存
        console.log('[green]已将解密文本保存至./text/Decryption_text_%s' %
                    (crypt_text_path.split('\\')[-1]))
        console.print('[#ffb74d]{Decryption 解密文本}：%s' % last_text)
    except:
        console.print_exception()
